﻿using System.Collections.Generic;

namespace theeg.MatrixWpf.Models
{
    public class SparseMatrix<TKey, TCell>
    {
        private readonly IDictionary<TKey, IDictionary<TKey, TCell>> _rows = new Dictionary<TKey, IDictionary<TKey, TCell>>();

        public void SetAt(TKey row, TKey col, TCell value)
        {
            if (Equals(value, default(TCell)))
            {
                RemoveAt(row, col);
            }
            else
            {
                IDictionary<TKey, TCell> storedRow;
                if (!_rows.TryGetValue(row, out storedRow))
                {
                    storedRow = new Dictionary<TKey, TCell>();
                    _rows.Add(row, storedRow);
                }
                storedRow[col] = value;
            }
        }

        public IDictionary<TKey, TCell> GetRow(TKey row)
        {
            IDictionary<TKey, TCell> storedRow;
            _rows.TryGetValue(row, out storedRow);
            return storedRow;
        }

        public TCell GetAt(TKey row, TKey col)
        {
            IDictionary<TKey, TCell> storedRow;
            if (_rows.TryGetValue(row, out storedRow))
            {
                TCell value;
                storedRow.TryGetValue(col, out value);
                return value;
            }
            return default(TCell);
        }

        public void RemoveAt(TKey row, TKey col)
        {
            IDictionary<TKey, TCell> storedRow;
            if (_rows.TryGetValue(row, out storedRow))
            {
                storedRow.Remove(col);
                if (storedRow.Count == 0)
                {
                    _rows.Remove(row);
                }
            }
        }
    }
}
